# import sys
from prettytable import PrettyTable
from sklearn.cluster import KMeans
from sklearn.decomposition import PCA
from sklearn.manifold import TSNE
from sklearn.pipeline import Pipeline
# from sko.AFSA import AFSA
from streamlit_extras.colored_header import colored_header

from business.algorithm.utils import *


def run():
    colored_header(label="机器学习：聚类与降维", description=" ", color_name="violet-90")
    file = st.file_uploader("Upload `.csv`file", type=['csv'], label_visibility="collapsed")
    if file is None:
        table = PrettyTable(['file name', 'class', 'description'])
        table.add_row(['file_1', 'dataset', 'data file'])
        st.write(table)
    if file is not None:
        df = pd.read_csv(file)
        check_string_NaN(df)
        colored_header(label="数据信息", description=" ", color_name="violet-70")
        nrow = st.slider("rows", 1, len(df), 5)
        df_nrow = df.head(nrow)
        st.write(df_nrow)

        colored_header(label="特征&目标", description=" ", color_name="violet-70")

        target_num = st.number_input('目标数量', min_value=1, max_value=10, value=1)

        col_feature, col_target = st.columns(2)

        # features
        features = df.iloc[:, :-target_num]
        # targets
        targets = df.iloc[:, -target_num:]
        with col_feature:
            st.write(features.head())
        with col_target:
            st.write(targets.head())
        cluster = CLUSTER(features, targets)

        colored_header(label="target", description=" ", color_name="violet-70")

        target_selected_name = st.selectbox('target', list(cluster.targets)[::-1])

        cluster.targets = targets[target_selected_name]
        # =============== cluster ================

        colored_header(label="聚类与降维", description=" ", color_name="violet-70")

        # colored_header(label="Choose Target", description=" ", color_name="violet-30")
        # target_selected_option = st.selectbox('target', list(cluster.targets)[::-1])

        # cluster.targets = targets[target_selected_option]

        model_path = './models/cluster'

        colored_header(label="Training", description=" ", color_name="violet-30")

        template_alg = model_platform(model_path)

        inputs, col2 = template_alg.show()

        if inputs['model'] == 'K-means':

            with col2:
                pass

            with st.container():
                button_train = st.button('Train', use_container_width=True)
            if button_train:
                cluster.model = KMeans(n_clusters=inputs['n clusters'], random_state=inputs['random state'])

                cluster.K_means()

                clustered_df = pd.concat([cluster.features, pd.DataFrame(cluster.model.labels_)], axis=1)

                r_name = 'cluster label'
                c_name = clustered_df.columns[-1]

                clustered_df.rename(columns={c_name: r_name}, inplace=True)
                with st.expander('cluster'):
                    st.write(clustered_df)
        if inputs['model'] == 'PCA':
            with col2:
                pass
            with st.container():
                button_train = st.button('Train', use_container_width=True)
            if button_train:
                pca_all = PCA(n_components=cluster.features.shape[1])
                pca_all.fit(cluster.features)
                with plt.style.context(['nature', 'no-latex']):
                    fig, ax = plt.subplots()
                    ax = plt.plot(np.cumsum(pca_all.explained_variance_ratio_ * 100))
                    plt.grid()
                    plt.xlabel('Numbers of components')
                    plt.ylabel('Explained variance')
                    st.pyplot(fig)

                def std_PCA(**argv):
                    scaler = MinMaxScaler()
                    pca = PCA(**argv)
                    pipeline = Pipeline([('scaler', scaler), ('pca', pca)])
                    return pipeline

                PCA_model = std_PCA(n_components=inputs['ncomponents'])
                PCA_transformed_data = PCA_model.fit_transform(cluster.features)
                if inputs['ncomponents'] == 2:
                    with plt.style.context(['nature', 'no-latex']):
                        fig, ax = plt.subplots()
                        ax = plt.scatter(PCA_transformed_data[:, 0], PCA_transformed_data[:, 1],
                                         c=[int(i) for i in cluster.targets.values], s=2, cmap='tab10')
                        plt.xlabel('1st dimension')
                        plt.ylabel('2st dimension')
                        plt.tight_layout()
                        st.pyplot(fig)
                    result_data = PCA_transformed_data
                    with st.expander('reduce dim'):
                        st.write(result_data)
                        tmp_download_link = download_button(result_data, f'dim reduction data.csv',
                                                            button_text='download')
                        st.markdown(tmp_download_link, unsafe_allow_html=True)
                else:
                    result_data = PCA_transformed_data
                    with st.expander('reduce dim'):
                        st.write(result_data)
                        tmp_download_link = download_button(result_data, f'dim reduction data.csv',
                                                            button_text='download')
                        st.markdown(tmp_download_link, unsafe_allow_html=True)
        if inputs['model'] == 'TSEN':
            with col2:
                pass
            with st.container():
                button_train = st.button('Train', use_container_width=True)
            if button_train:
                TSNE_model = TSNE(n_components=inputs['ncomponents'], perplexity=inputs['perplexity'],
                                  learning_rate='auto', n_iter=inputs['max iter'], init='pca',
                                  random_state=inputs['random state'])
                TSNE_transformed_data = TSNE_model.fit_transform(cluster.features)

                if inputs['ncomponents'] == 2:
                    with plt.style.context(['nature', 'no-latex']):
                        fig, ax = plt.subplots()
                        ax = plt.scatter(TSNE_transformed_data[:, 0], TSNE_transformed_data[:, 1],
                                         c=[int(i) for i in cluster.targets.values], s=2, cmap='tab10')
                        plt.xlabel('1st dimension')
                        plt.ylabel('2st dimension')
                        plt.tight_layout()
                        st.pyplot(fig)
                    result_data = TSNE_transformed_data
                    with st.expander('reduce dim'):
                        st.write(result_data)
                        tmp_download_link = download_button(result_data, f'dim reduction data.csv',
                                                            button_text='download')
                        st.markdown(tmp_download_link, unsafe_allow_html=True)
                else:
                    result_data = TSNE_transformed_data
                    with st.expander('reduce dim'):
                        st.write(result_data)
                        tmp_download_link = download_button(result_data, f'dim reduction data.csv',
                                                            button_text='download')
                        st.markdown(tmp_download_link, unsafe_allow_html=True)
        st.write('---')